%% ========================================================================
%               调度方案约束验证与目标值计算脚本 (V4)
% ========================================================================
%
% V4 更新:
% - [Bug修复] 使用 ismember 函数替换了 || 运算符，以处理 code_map 中
%   可能存在的空值，避免"操作数必须为标量逻辑值"的错误。
%
% V3 更新:
% - [核心修正] 增加了对L0缓存的驻留数量检查。
% - [核心修正] 重写了maxVstay的计算逻辑，精确地只统计L1和UB。
%
% V2 更新:
% - [Bug修复] 对加载的schedule进行了+1操作以修正索引。
%
clear; clc; close all;

%% ========================================================================
%                            第一部分: 设置
% ========================================================================
% --- 请在这里指定要验证的解文件名 ---
solution_to_validate_filename = '../c_implementation/result-all/FlashAttention_Case1/FlashAttention_Case1_schedule.txt'; % <--- 验证FlashAttention_Case1结果

%% ========================================================================
%             第二部分: 加载数据与预计算 (与main.m保持一致)
% ========================================================================
disp('--- 正在加载验证所需的数据和环境... ---');
try
    nodes_table = readtable('../c_implementation/data/CSV-version/FlashAttention_Case1_Nodes.csv');  % 验证FlashAttention_Case1
    edges_table = readtable('../c_implementation/data/CSV-version/FlashAttention_Case1_Edges.csv');  % 验证FlashAttention_Case1
catch ME
    disp('错误：无法加载CSV文件。请检查文件路径是否正确。'); rethrow(ME);
end

% ... (与main.m中相同的预处理代码) ...
nodes_table.OpCode = double(categorical(nodes_table.Op));
nodes_table.TypeCode = double(categorical(nodes_table.Type));
categories_op = categories(categorical(nodes_table.Op));
categories_type = categories(categorical(nodes_table.Type));
code_map.op.ALLOC = find(strcmp('ALLOC', categories_op), 1);
code_map.op.FREE = find(strcmp('FREE', categories_op), 1);
code_map.type.L1 = find(strcmp('L1', categories_type), 1);
code_map.type.UB = find(strcmp('UB', categories_type), 1);
code_map.type.L0A = find(strcmp('L0A', categories_type), 1);
code_map.type.L0B = find(strcmp('L0B', categories_type), 1);
code_map.type.L0C = find(strcmp('L0C', categories_type), 1);
G = digraph(edges_table.StartNodeId + 1, edges_table.EndNodeId + 1, [], nodes_table);
num_nodes = numnodes(G);

% 预计算节点间的依赖关系
is_ancestor = full(adjacency(transclosure(G)));
disp('--- 数据加载与预计算完成 ---');

%% ========================================================================
%                       第三部分: 加载待验证的解
% ========================================================================
fprintf('\n--- 正在从 "%s" 加载调度方案... ---\n', solution_to_validate_filename);
try
    schedule_to_validate = readmatrix(solution_to_validate_filename);
    schedule_to_validate = schedule_to_validate + 1; 
    if size(schedule_to_validate, 2) > 1, schedule_to_validate = schedule_to_validate'; end
    fprintf('成功加载并转换调度方案，包含 %d 个步骤。\n', length(schedule_to_validate));
catch ME
    fprintf('错误：无法读取文件 "%s"。请检查文件名和文件内容。\n', solution_to_validate_filename);
    rethrow(ME);
end

%% ========================================================================
%                          第四部分: 约束验证
% ========================================================================
disp('--- 开始验证约束条件... ---');
validation_passed = true;

% --- 检查 1: 完整性检查 ---
if numel(unique(schedule_to_validate)) ~= num_nodes || numel(schedule_to_validate) ~= num_nodes
    fprintf('[验证失败] 约束: 完整性约束\n');
    fprintf('  > 原因: 调度方案未包含所有节点，或存在重复/缺失节点。\n');
    validation_passed = false;
end

% --- 检查 2: 依赖约束 (拓扑排序) ---
if validation_passed
    node_positions = zeros(num_nodes, 1);
    node_positions(schedule_to_validate) = 1:num_nodes;
    is_dependency_valid = true;
    for u = 1:num_nodes
        for v = 1:num_nodes
            if is_ancestor(u, v) && (node_positions(u) >= node_positions(v))
                fprintf('[验证失败] 约束: 依赖约束 (拓扑排序)\n');
                fprintf('  > 原因: 节点 %d (祖先) 被安排在了节点 %d (后代) 的后面。\n', u-1, v-1);
                is_dependency_valid = false; break;
            end
        end
        if ~is_dependency_valid, break; end
    end
    if ~is_dependency_valid, validation_passed = false; else, disp('[验证通过] 约束: 依赖约束 (拓扑排序)'); end
end

% --- 检查 3: L0 缓存驻留约束 ---
if validation_passed
    l0_counts = containers.Map({'L0A', 'L0B', 'L0C'}, {0, 0, 0});
    l0_type_codes = [code_map.type.L0A, code_map.type.L0B, code_map.type.L0C];
    l0_type_names = {'L0A', 'L0B', 'L0C'};
    is_l0_valid = true;
    
    for i = 1:length(schedule_to_validate)
        node_id = schedule_to_validate(i);
        node_info = G.Nodes(node_id, :);
        node_op_code = node_info.OpCode;
        node_type_code = node_info.TypeCode;
        
        [is_l0, type_idx] = ismember(node_type_code, l0_type_codes);
        if is_l0
            type_name = l0_type_names{type_idx};
            if node_op_code == code_map.op.ALLOC
                if l0_counts(type_name) > 0
                    fprintf('[验证失败] 约束: L0缓存驻留约束\n');
                    fprintf('  > 原因: 在执行步骤 %d (节点ID %d) 时，试图申请 %s 缓存，但该缓存上已有缓冲区驻留。\n', i, node_id-1, type_name);
                    is_l0_valid = false; break;
                end
                l0_counts(type_name) = l0_counts(type_name) + 1;
            elseif node_op_code == code_map.op.FREE
                l0_counts(type_name) = l0_counts(type_name) - 1;
            end
        end
    end
    if ~is_l0_valid, validation_passed = false; else, disp('[验证通过] 约束: L0缓存驻留约束'); end
end


%% ========================================================================
%                      第五部分: 计算目标值
% ========================================================================
if validation_passed
    fprintf('\n--- 所有约束验证通过! 开始计算目标值 (仅L1和UB)... ---\n');
    
    % --- 内置 maxVstay 计算逻辑 ---
    v_stay = 0;
    max_v_stay = 0;
    
    for i = 1:length(schedule_to_validate)
        node_id = schedule_to_validate(i);
        node_info = G.Nodes(node_id, :);
        node_type_code = node_info.TypeCode;
        
        % --- 【Bug修复】使用 ismember 替换 || 来处理可能的空值 ---
        % 只考虑 L1 和 UB 类型的缓存
        if ismember(node_type_code, [code_map.type.L1, code_map.type.UB])
            node_op_code = node_info.OpCode;
            node_size = node_info.Size;
            if node_op_code == code_map.op.ALLOC
                v_stay = v_stay + node_size;
            elseif node_op_code == code_map.op.FREE
                v_stay = v_stay - node_size;
            end
            if v_stay > max_v_stay
                max_v_stay = v_stay;
            end
        end
    end
    
    fprintf('\n==================================\n');
    fprintf('  最终结果\n');
    fprintf('----------------------------------\n');
    fprintf('  调度方案: %s\n', solution_to_validate_filename);
    fprintf('  状态:     有效 (Valid)\n');
    fprintf('  目标值 (maxVstay for L1+UB): %d\n', max_v_stay);
    fprintf('==================================\n');

else
    fprintf('\n==================================\n');
    fprintf('  最终结果\n');
    fprintf('----------------------------------\n');
    fprintf('  调度方案: %s\n', solution_to_validate_filename);
    fprintf('  状态:     无效 (Invalid)\n');
    fprintf('  > 请检查上方日志中的失败原因。\n');
    fprintf('==================================\n');
end